﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.IO;
using System.Management;
using System.Runtime.InteropServices;
using System.Windows.Forms;

namespace U盘格式工具
{
    public partial class MainForm : Form
    {
        private List<string> _filesToCopy = new List<string>();
        private string _selectedFilesInfo = "未选择文件";
        private const int WM_DEVICECHANGE = 0x219;
        private const int DBT_DEVICEARRIVAL = 0x8000;
        private const int DBT_DEVICEREMOVECOMPLETE = 0x8004;
        private const int DBT_DEVTYP_VOLUME = 0x00000002;

        public MainForm()
        {
            InitializeComponent();
            RefreshUSBDevices();
            UpdateSelectedFilesDisplay();

            this.AllowDrop = true;
            this.DragEnter += MainForm_DragEnter;
            this.DragDrop += MainForm_DragDrop;
        }

        protected override void WndProc(ref Message m)
        {
            base.WndProc(ref m);

            if (m.Msg == WM_DEVICECHANGE)
            {
                int eventType = m.WParam.ToInt32();
                if (eventType == DBT_DEVICEARRIVAL || eventType == DBT_DEVICEREMOVECOMPLETE)
                {
                    DEV_BROADCAST_HDR hdr = (DEV_BROADCAST_HDR)Marshal.PtrToStructure(m.LParam, typeof(DEV_BROADCAST_HDR));

                    if (hdr.dbch_devicetype == DBT_DEVTYP_VOLUME)
                    {
                        this.BeginInvoke(new MethodInvoker(() =>
                        {
                            RefreshUSBDevices();
                            txtLog.AppendText($"检测到U盘{(eventType == DBT_DEVICEARRIVAL ? "插入" : "拔出")}，已刷新列表\n");
                        }));
                    }
                }
            }
        }

        [StructLayout(LayoutKind.Sequential)]
        private struct DEV_BROADCAST_HDR
        {
            public int dbch_size;
            public int dbch_devicetype;
            public int dbch_reserved;
        }

        private void MainForm_DragEnter(object sender, DragEventArgs e)
        {
            if (e.Data.GetDataPresent(DataFormats.FileDrop))
                e.Effect = DragDropEffects.Copy;
            else
                e.Effect = DragDropEffects.None;
        }

        private void btnEject_Click(object sender, EventArgs e)
        {
            var selectedDrives = GetSelectedDriveLetters();
            if (selectedDrives.Count == 0)
            {
                MessageBox.Show("请先选择要弹出的U盘设备", "提示", MessageBoxButtons.OK, MessageBoxIcon.Information);
                return;
            }

            btnEject.Enabled = false;
            btnRefresh.Enabled = false;
            btnFormatSelected.Enabled = false;

            foreach (var drive in selectedDrives)
            {
                txtLog.AppendText($"正在弹出 {drive}...\n");
                EjectDriveUsingShell(drive);
            }

            btnEject.Enabled = true;
            btnRefresh.Enabled = true;
            btnFormatSelected.Enabled = true;
            RefreshUSBDevices();
        }

        private bool EjectDriveUsingShell(string driveLetter)
        {
            try
            {
                Type shellAppType = Type.GetTypeFromProgID("Shell.Application");
                dynamic shell = Activator.CreateInstance(shellAppType);

                foreach (var item in shell.NameSpace(17).Items())
                {
                    string drivePath = item.Path;
                    if (drivePath.StartsWith(driveLetter, StringComparison.OrdinalIgnoreCase))
                    {
                        item.InvokeVerb("Eject");
                        txtLog.AppendText($"已请求弹出 {driveLetter}\n");
                        return true;
                    }
                }

                txtLog.AppendText($"未找到驱动器 {driveLetter}\n");
                return false;
            }
            catch (Exception ex)
            {
                txtLog.AppendText($"弹出 {driveLetter} 失败: {ex.Message}\n");
                return false;
            }
        }

        private void MainForm_DragDrop(object sender, DragEventArgs e)
        {
            HandleDroppedFiles(e);
        }

        private void HandleDroppedFiles(DragEventArgs e)
        {
            _filesToCopy.Clear();
            _selectedFilesInfo = "";

            if (e.Data.GetDataPresent(DataFormats.FileDrop))
            {
                string[] droppedItems = (string[])e.Data.GetData(DataFormats.FileDrop);
                int fileCount = 0;
                int folderCount = 0;

                foreach (string item in droppedItems)
                {
                    if (File.Exists(item))
                    {
                        _filesToCopy.Add(item);
                        fileCount++;
                        _selectedFilesInfo += $"文件: {Path.GetFileName(item)}\n";
                    }
                    else if (Directory.Exists(item))
                    {
                        _filesToCopy.Add(item);
                        folderCount++;
                        _selectedFilesInfo += $"文件夹: {Path.GetFileName(item)}\n";
                    }
                }

                _selectedFilesInfo = $"已拖放 {fileCount} 个文件和 {folderCount} 个文件夹:\n" + _selectedFilesInfo;
                UpdateSelectedFilesDisplay();
            }
        }

        private void RefreshUSBDevices()
        {
            lstUSBDrives.Items.Clear();

            try
            {
                var searcher = new ManagementObjectSearcher(
                    "SELECT * FROM Win32_DiskDrive WHERE InterfaceType='USB'");

                bool hasDevices = false;
                foreach (ManagementObject disk in searcher.Get())
                {
                    hasDevices = true;
                    var partitionQueryText = string.Format(
                        "ASSOCIATORS OF {{Win32_DiskDrive.DeviceID='{0}'}} WHERE AssocClass = Win32_DiskDriveToDiskPartition",
                        disk["DeviceID"]);

                    var partitionSearcher = new ManagementObjectSearcher(partitionQueryText);
                    foreach (ManagementObject partition in partitionSearcher.Get())
                    {
                        var logicalDiskQueryText = string.Format(
                            "ASSOCIATORS OF {{Win32_DiskPartition.DeviceID='{0}'}} WHERE AssocClass = Win32_LogicalDiskToPartition",
                            partition["DeviceID"]);

                        var logicalDiskSearcher = new ManagementObjectSearcher(logicalDiskQueryText);
                        foreach (ManagementObject logicalDisk in logicalDiskSearcher.Get())
                        {
                            string driveLetter = logicalDisk["DeviceID"].ToString();
                            string volumeName = logicalDisk["VolumeName"]?.ToString() ?? "无标签";
                            string fileSystem = logicalDisk["FileSystem"]?.ToString() ?? "未知";
                            decimal size = Convert.ToDecimal(logicalDisk["Size"]);
                            decimal freeSpace = Convert.ToDecimal(logicalDisk["FreeSpace"]);

                            var item = new ListViewItem(new[] {
                                driveLetter,
                                volumeName,
                                fileSystem,
                                FormatSize(size),
                                FormatSize(freeSpace),
                                disk["Model"].ToString()
                            });
                            item.Tag = driveLetter;
                            lstUSBDrives.Items.Add(item);
                        }
                    }
                }

                if (!hasDevices)
                {
                    var item = new ListViewItem("没有检测到USB设备");
                    item.Checked = false;
                    item.Tag = null;
                    lstUSBDrives.Items.Add(item);
                }
            }
            catch (Exception ex)
            {
                MessageBox.Show($"查询USB设备时出错: {ex.Message}", "错误", MessageBoxButtons.OK, MessageBoxIcon.Error);
                var errorItem = new ListViewItem("查询USB设备失败");
                errorItem.Checked = false;
                errorItem.Tag = null;
                lstUSBDrives.Items.Add(errorItem);
            }
        }

        private void btnSelectAll_Click(object sender, EventArgs e)
        {
            bool allSelected = true;

            foreach (ListViewItem item in lstUSBDrives.Items)
            {
                if (!item.Checked)
                {
                    allSelected = false;
                    break;
                }
            }

            foreach (ListViewItem item in lstUSBDrives.Items)
            {
                item.Checked = !allSelected;
            }

            btnSelectAll.Text = allSelected ? "全选" : "取消全选";
        }

        private string FormatSize(decimal bytes)
        {
            if (bytes <= 0) return "0 B";

            string[] suffixes = { "B", "KB", "MB", "GB", "TB" };
            int suffixIndex = 0;
            decimal size = bytes;

            while (size >= 1024 && suffixIndex < suffixes.Length - 1)
            {
                size /= 1024;
                suffixIndex++;
            }

            return $"{size:n2} {suffixes[suffixIndex]}";
        }

        private void btnRefresh_Click(object sender, EventArgs e)
        {
            RefreshUSBDevices();
        }

        private void btnFormatSelected_Click(object sender, EventArgs e)
        {
            var selectedDrives = GetSelectedDriveLetters();
            if (selectedDrives.Count == 0)
            {
                MessageBox.Show("请先选择有效的U盘设备", "提示", MessageBoxButtons.OK, MessageBoxIcon.Information);
                return;
            }

            string fileSystem = cmbFileSystem.SelectedItem?.ToString() ?? "FAT32";
            string volumeLabel = txtVolumeLabel.Text.Trim();
            bool quickFormat = chkQuickFormat.Checked;
            string targetFolder = txtTargetFolder.Text.Trim();

            var confirmResult = MessageBox.Show(
                $"确定要格式化选中的 {lstUSBDrives.CheckedItems.Count} 个U盘吗?\n" +
                $"文件系统: {fileSystem}\n卷标: {volumeLabel}\n" +
                $"快速格式化: {quickFormat}\n" +
                (_filesToCopy.Count > 0 ? $"将复制 {_filesToCopy.Count} 个文件/文件夹到: {targetFolder}" : "不复制文件"),
                "确认格式化",
                MessageBoxButtons.YesNo,
                MessageBoxIcon.Warning);

            if (confirmResult != DialogResult.Yes)
            {
                return;
            }

            btnFormatSelected.Enabled = false;
            btnRefresh.Enabled = false;
            btnSelectFiles.Enabled = false;
            btnSelectFolder.Enabled = false;

            backgroundWorker1.RunWorkerAsync(new FormatArgs
            {
                DriveLetters = GetSelectedDriveLetters(),
                FileSystem = fileSystem,
                VolumeLabel = volumeLabel,
                QuickFormat = quickFormat,
                FilesToCopy = _filesToCopy,
                TargetFolder = targetFolder
            });
        }

        private List<string> GetSelectedDriveLetters()
        {
            var drives = new List<string>();
            foreach (ListViewItem item in lstUSBDrives.CheckedItems)
            {
                if (item.Tag != null)
                {
                    drives.Add(item.Tag.ToString());
                }
            }
            return drives;
        }

        private void backgroundWorker1_DoWork(object sender, DoWorkEventArgs e)
        {
            var args = (FormatArgs)e.Argument;
            int completed = 0;

            foreach (var drive in args.DriveLetters)
            {
                if (backgroundWorker1.CancellationPending)
                {
                    e.Cancel = true;
                    return;
                }

                try
                {
                    string drivePath = drive.Length == 1 ? drive + ":" : drive;
                    drivePath = drivePath.Substring(0, 2);

                    backgroundWorker1.ReportProgress(0, $"正在格式化 {drivePath}...");

                    var psi = new ProcessStartInfo
                    {
                        FileName = "format.com",
                        Arguments = $"{drivePath} /FS:{args.FileSystem} /V:{args.VolumeLabel} {(args.QuickFormat ? "/Q" : "")} /Y",
                        UseShellExecute = false,
                        CreateNoWindow = true,
                        RedirectStandardInput = true,
                        RedirectStandardOutput = true,
                        RedirectStandardError = true
                    };

                    using (var process = new Process { StartInfo = psi })
                    {
                        process.Start();
                        process.StandardInput.WriteLine();

                        string output = process.StandardOutput.ReadToEnd();
                        string error = process.StandardError.ReadToEnd();

                        process.WaitForExit();

                        if (process.ExitCode != 0)
                        {
                            backgroundWorker1.ReportProgress(0, $"格式化 {drivePath} 失败: {error}");
                            continue;
                        }
                        backgroundWorker1.ReportProgress(0, $"成功格式化 {drivePath}");
                    }

                    if (args.FilesToCopy != null && args.FilesToCopy.Count > 0)
                    {
                        System.Threading.Thread.Sleep(args.RemountDelay);

                        string targetPath = Path.Combine(drivePath, args.TargetFolder.TrimStart('\\'));

                        backgroundWorker1.ReportProgress(0, $"正在复制文件到 {targetPath}...");

                        if (!Directory.Exists(targetPath))
                        {
                            Directory.CreateDirectory(targetPath);
                        }

                        foreach (var sourcePath in args.FilesToCopy)
                        {
                            if (File.Exists(sourcePath))
                            {
                                string destFile = Path.Combine(targetPath, Path.GetFileName(sourcePath));
                                File.Copy(sourcePath, destFile, true);
                                backgroundWorker1.ReportProgress(0, $"已复制文件: {Path.GetFileName(sourcePath)}");
                            }
                            else if (Directory.Exists(sourcePath))
                            {
                                string destDir = Path.Combine(targetPath, Path.GetFileName(sourcePath));
                                CopyDirectory(sourcePath, destDir);
                                backgroundWorker1.ReportProgress(0, $"已复制文件夹: {Path.GetFileName(sourcePath)}");
                            }
                        }

                        backgroundWorker1.ReportProgress(0, $"文件复制到 {drivePath} 完成");
                    }
                }
                catch (Exception ex)
                {
                    backgroundWorker1.ReportProgress(0, $"处理 {drive} 时出错: {ex.Message}");
                }

                completed++;
                int progress = (int)((float)completed / args.DriveLetters.Count * 100);
                backgroundWorker1.ReportProgress(progress);
            }
        }

        private void CopyDirectory(string sourceDir, string destinationDir)
        {
            var dir = new DirectoryInfo(sourceDir);

            if (!dir.Exists)
                throw new DirectoryNotFoundException($"源目录不存在: {sourceDir}");

            Directory.CreateDirectory(destinationDir);

            foreach (FileInfo file in dir.GetFiles())
            {
                string targetFilePath = Path.Combine(destinationDir, file.Name);
                file.CopyTo(targetFilePath, true);
            }

            foreach (DirectoryInfo subDir in dir.GetDirectories())
            {
                string newDestinationDir = Path.Combine(destinationDir, subDir.Name);
                CopyDirectory(subDir.FullName, newDestinationDir);
            }
        }

        private void backgroundWorker1_ProgressChanged(object sender, ProgressChangedEventArgs e)
        {
            if (e.UserState != null)
            {
                txtLog.AppendText(e.UserState.ToString() + Environment.NewLine);
            }

            progressBar1.Value = e.ProgressPercentage;
        }

        private void backgroundWorker1_RunWorkerCompleted(object sender, RunWorkerCompletedEventArgs e)
        {
            btnFormatSelected.Enabled = true;
            btnRefresh.Enabled = true;
            btnSelectFiles.Enabled = true;
            btnSelectFolder.Enabled = true;

            if (e.Cancelled)
            {
                txtLog.AppendText("操作已取消" + Environment.NewLine);
            }
            else if (e.Error != null)
            {
                txtLog.AppendText($"发生错误: {e.Error.Message}" + Environment.NewLine);
            }
            else
            {
                txtLog.AppendText("所有操作已完成" + Environment.NewLine);
                RefreshUSBDevices();
            }
        }

        private void MainForm_Load(object sender, EventArgs e)
        {
            cmbFileSystem.Items.AddRange(new[] { "FAT32", "NTFS", "exFAT" });
            cmbFileSystem.SelectedIndex = 0;

            lstUSBDrives.View = View.Details;
            lstUSBDrives.CheckBoxes = true;
            lstUSBDrives.FullRowSelect = true;
            lstUSBDrives.Columns.Add("驱动器", 80);
            lstUSBDrives.Columns.Add("卷标", 100);
            lstUSBDrives.Columns.Add("文件系统", 80);
            lstUSBDrives.Columns.Add("容量", 100);
            lstUSBDrives.Columns.Add("可用空间", 100);
            lstUSBDrives.Columns.Add("型号", 200);
        }

        private void btnSelectFiles_Click(object sender, EventArgs e)
        {
            using (var dialog = new OpenFileDialog())
            {
                dialog.Multiselect = true;
                dialog.Title = "选择要复制到U盘的文件";
                if (dialog.ShowDialog() == DialogResult.OK)
                {
                    _filesToCopy.Clear();
                    _filesToCopy.AddRange(dialog.FileNames);
                    _selectedFilesInfo = $"已选择 {_filesToCopy.Count} 个文件:\n";
                    foreach (var file in _filesToCopy)
                    {
                        _selectedFilesInfo += $"- {file}\n";
                    }
                    UpdateSelectedFilesDisplay();
                }
            }
        }

        private void btnSelectFolder_Click(object sender, EventArgs e)
        {
            using (var dialog = new FolderBrowserDialog())
            {
                dialog.Description = "选择要复制到U盘的文件夹";
                if (dialog.ShowDialog() == DialogResult.OK)
                {
                    _filesToCopy.Clear();
                    _filesToCopy.Add(dialog.SelectedPath);
                    _selectedFilesInfo = $"已选择文件夹:\n- {dialog.SelectedPath}";
                    UpdateSelectedFilesDisplay();
                }
            }
        }

        private void UpdateSelectedFilesDisplay()
        {
            txtSelectedFiles.Text = _selectedFilesInfo;
            txtSelectedFiles.ScrollBars = ScrollBars.Vertical;
        }

        private void btnClearSelection_Click(object sender, EventArgs e)
        {
            _filesToCopy.Clear();
            _selectedFilesInfo = "未选择文件";
            UpdateSelectedFilesDisplay();
        }
    }

    public class FormatArgs
    {
        public List<string> DriveLetters { get; set; }
        public string FileSystem { get; set; }
        public string VolumeLabel { get; set; }
        public bool QuickFormat { get; set; }
        public List<string> FilesToCopy { get; set; }
        public string TargetFolder { get; set; }
        public int RemountDelay { get; set; } = 1500;
    }
}